from fastai.vision import *

from .filters import IFilter, MasterFilter, ColorizerFilter
from .generators import gen_inference_wide


class ModelImageVisualizer:
    def __init__(self, filter: IFilter, results_dir: str = None):
        self.filter = filter
        self.results_dir = None if results_dir is None else Path(results_dir)
        self.results_dir.mkdir(parents=True, exist_ok=True)

    def _clean_mem(self):
        torch.cuda.empty_cache()
        # gc.collect()

    def get_transformed_image(
            self, orig_image, render_factor: int = None
    ) -> Image:
        self._clean_mem()
        filtered_image = self.filter.filter(
            orig_image, orig_image, render_factor=render_factor
        )
        return filtered_image


def get_image_colorizer(
        render_factor: int = 35
) -> ModelImageVisualizer:
    return get_stable_image_colorizer(render_factor=render_factor)


def get_stable_image_colorizer(
        root_folder: Path = Path('./'),
        weights_name: str = 'ColorizeStable_gen',
        results_dir='result_images',
        render_factor: int = 35,
) -> ModelImageVisualizer:
    learn = gen_inference_wide(root_folder=root_folder, weights_name=weights_name)
    filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)
    vis = ModelImageVisualizer(filtr, results_dir=results_dir)
    return vis
